module BTree
  (
    BTree(..),
    getNodeWeight,
    getNodeChar,
    getNodePath,
    walkHuffmanTree,
    walkLeaf
  ) where

data BTree = Node {
  -- string info for huffman code
  nodeInfo :: (Int, Char, String),
  leftTree :: BTree,
  rightTree :: BTree
                  } | Empty deriving (Show, Eq)

getNodeWeight :: BTree -> Int
getNodeWeight tree = case nodeInfo tree of
                       (weight, _, _) -> weight

getNodeChar :: BTree -> Char
getNodeChar tree = case nodeInfo tree of
                     (_, ch, _) -> ch

getNodePath :: BTree -> String
getNodePath tree = case nodeInfo tree of
                     (_, _, path) -> path

walkOnce :: BTree -> [BTree]
walkOnce Empty = []
walkOnce (Node c leftTree rightTree) =
  (Node c leftTree rightTree) : walkOnce leftTree

walkLeaf :: BTree -> [(Int, Char, String)]
walkLeaf (Node c Empty Empty) = [c]
walkLeaf (Node c leftTree rightTree) =
  (walkLeaf leftTree) ++ (walkLeaf rightTree)

walkHuffmanTree :: String -> BTree -> [(Char, String)]
walkHuffmanTree curPath (Node (_, ch, _) Empty Empty) = [(ch, curPath)]
walkHuffmanTree curPath (Node info lTree rTree)=
  (walkHuffmanTree (curPath++"0") lTree) ++ (walkHuffmanTree (curPath++"1") rTree)

getTreeDepth :: BTree -> Int
getTreeDepth Empty = 0
getTreeDepth (Node c leftTree rightTree) =
  max (getTreeDepth leftTree) (getTreeDepth rightTree) + 1
    where max a b = if a >= b
                       then a
                       else b
